import argparse
import tvm
# import tvm.relay as relay
from tvm import relay
import torchvision
import torch
from keras_model_impl import LeNet


def run(args):
    # model = tf.keras.Sequential()
    # model.add(tf.keras.layers.Conv2D(32, (3, 3), activation=args.at, input_shape=(28, 28, 1)))

    # inputs1 = layers.Input(shape=(224, 224, 3), batch_size=1, name="input_1")
    # x = layers.Conv2D(32, 3)(inputs1)
    # x = layers.BatchNormalization()(inputs1)
    # model = Model(inputs=inputs1, outputs=x)

    # onnx_model = keras2onnx.convert_keras(model, model.name)
    # onnx.save(onnx_model, "{}.onnx".format(args.output))
    # node = onnx_model.graph.node[0]

    model_name = "torch-resnet18"
    image_size = 224
    shape_dict = [("input0", (1, 3, image_size, image_size))]
    # shape_dict = [("input0", (1, image_size))]
    # model = LeNet()
    model = torchvision.models.mobilenet_v2()
    model.eval()

    # model = torch.hub.load('facebookresearch/pytorch_GAN_zoo:hub', 'DCGAN', pretrained=True)
    # model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch',
    #                        'unet',
    #                        in_channels=3,
    #                        out_channels=1,
    #                        init_features=32,
    #                        pretrained=True)
    # model = torchvision.models.mnasnet1_0()
    # model = torch.hub.load('pytorch/vision:v0.10.0', 'fcn_resnet50', pretrained=True)

    input_shape = list(shape_dict[0][1])

    input_data = torch.randn(input_shape)
    scripted_model = torch.jit.trace(model, input_data).eval()
    # model = keras.applications.VGG19()
    mod, params = relay.frontend.from_pytorch(scripted_model, shape_dict)
    # exit(0)
    with tvm.transform.PassContext(opt_level=args.opt_level):
        lib = relay.build(mod, "llvm -mcpu=skylake-avx512", params=params)
    lib.export_library("{}.so".format(model_name))

    json_str = lib["get_json"]()
    with open("{}.json".format(model_name), "w") as f:
        f.write(json_str)

    return model


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='TVM Reversion')
    parser.add_argument("--at", default="relu", type=str)  # activation type
    parser.add_argument("-o", "--output", default="out.so", type=str)
    parser.add_argument("-ol", "--opt_level", default=1, type=int)
    parser.set_defaults(show_failure_case=False)
    args = parser.parse_args()
    run(args)
